package tree.structure.trie;

/**
 * 前缀树
 * - 根节点不包含字符，除根节点外每一个节点都只包含一个字符
 * - 从根节点到某一节点，路径上经过的字符连接起来，为该节点对应的字符串
 * - 每个节点的所有子节点包含的字符都不相同
 */
public class Trie {

    private TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    public void insert(String s) {
        // 已加入前缀树的字符串个数
        root.addPass(1);
        char[] arr = s.toCharArray();

        TrieNode tempRoot = root;
        for (char c : arr) {
            int idx = c - 'a';
            if (tempRoot.getTrieNodes()[idx] == null) {
                tempRoot.getTrieNodes()[idx] = new TrieNode();
            }
            tempRoot.getTrieNodes()[idx].addPass(1);
            tempRoot = tempRoot.getTrieNodes()[idx];
        }
        // 该处为一个字符串的结尾
        tempRoot.addEnd(1);
    }

    public int countWord(String s) {
        char[] arr = s.toCharArray();
        TrieNode tempRoot = root;
        for (char ch : arr) {
            int idx = ch - 'a';
            if (tempRoot.getTrieNodes()[idx] == null) {
                return 0;
            }
            tempRoot = tempRoot.getTrieNodes()[idx];
        }
        return tempRoot.getEnd();
    }

    public int countPrefixWord(String s) {
        char[] arr = s.toCharArray();
        TrieNode tempRoot = root;

        for (char ch : arr) {
            int idx = ch - 'a';
            if (tempRoot.getTrieNodes()[idx] == null) {
                return 0;
            }
            tempRoot = tempRoot.getTrieNodes()[idx];
        }
        return tempRoot.getPass();
    }

    public void erase(String s) {
        if (countWord(s) < 1) {
            return;
        }
        char[] arr = s.toCharArray();
        TrieNode tempRoot = root;
        tempRoot.addPass(-1);

        for (char ch : arr) {
            int idx = ch - 'a';

            TrieNode trieNode = tempRoot.getTrieNodes()[idx];
            trieNode.addPass(-1);

            if (trieNode.getPass() == 0) {
                tempRoot.getTrieNodes()[idx] = null;
                return;
            }
            tempRoot = trieNode;
        }

        tempRoot.addEnd(-1);
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("a");
        trie.insert("abc");
        trie.insert("ad");
        trie.insert("ab");
        trie.insert("ab");

        trie.erase("ab");

        System.out.println("trie.countWord(\"a\") = " + trie.countWord("a"));
        System.out.println("trie.countWord(\"aa\") = " + trie.countWord("aa"));
        System.out.println("trie.countWord(\"ab\") = " + trie.countWord("ab"));

        System.out.println("trie = " + trie);
    }

}
